home *** CD-ROM | disk | FTP | other *** search
/ PC World Komputer 2010 April / PCWorld0410.iso / hity wydania / Ubuntu 9.10 PL / karmelkowy-koliberek-9.10-netbook-remix-PL.iso / casper / filesystem.squashfs / usr / share / pyshared / couchdb / multipart.py < prev    next >
Text File  |  2009-07-02  |  9KB  |  248 lines
  1. # -*- coding: utf-8 -*-
  2. #
  3. # Copyright (C) 2008-2009 Christopher Lenz
  4. # All rights reserved.
  5. #
  6. # This software is licensed as described in the file COPYING, which
  7. # you should have received as part of this distribution.
  8.  
  9. """Support for streamed reading and writing of multipart MIME content."""
  10.  
  11. from base64 import b64encode
  12. from cgi import parse_header
  13. try:
  14.     from hashlib import md5
  15. except ImportError:
  16.     from md5 import new as md5
  17. import sys
  18.  
  19. __all__ = ['read_multipart', 'write_multipart']
  20. __docformat__ = 'restructuredtext en'
  21.  
  22.  
  23. CRLF = '\r\n'
  24.  
  25.  
  26. def read_multipart(fileobj, boundary=None):
  27.     """Simple streaming MIME multipart parser.
  28.     
  29.     This function takes a file-like object reading a MIME envelope, and yields
  30.     a ``(headers, is_multipart, payload)`` tuple for every part found, where
  31.     ``headers`` is a dictionary containing the MIME headers of that part (with
  32.     names lower-cased), ``is_multipart`` is a boolean indicating whether the
  33.     part is itself multipart, and ``payload`` is either a string (if
  34.     ``is_multipart`` is false), or an iterator over the nested parts.
  35.     
  36.     Note that the iterator produced for nested multipart payloads MUST be fully
  37.     consumed, even if you wish to skip over the content.
  38.     
  39.     :param fileobj: a file-like object
  40.     :param boundary: the part boundary string, will generally be determined
  41.                      automatically from the headers of the outermost multipart
  42.                      envelope
  43.     :return: an iterator over the parts
  44.     :since: 0.5
  45.     """
  46.     headers = {}
  47.     buf = []
  48.     outer = in_headers = boundary is None
  49.  
  50.     next_boundary = boundary and '--' + boundary + '\n' or None
  51.     last_boundary = boundary and '--' + boundary + '--\n' or None
  52.  
  53.     def _current_part():
  54.         payload = ''.join(buf)
  55.         if payload.endswith('\r\n'):
  56.             payload = payload[:-2]
  57.         elif payload.endswith('\n'):
  58.             payload = payload[:-1]
  59.         content_md5 = headers.get('content-md5')
  60.         if content_md5:
  61.             h = b64encode(md5(payload).digest())
  62.             if content_md5 != h:
  63.                 raise ValueError('data integrity check failed')
  64.         return headers, False, payload
  65.  
  66.     for line in fileobj:
  67.         if in_headers:
  68.             line = line.replace(CRLF, '\n')
  69.             if line != '\n':
  70.                 name, value = line.split(':', 1)
  71.                 headers[name.lower().strip()] = value.strip()
  72.             else:
  73.                 in_headers = False
  74.                 mimetype, params = parse_header(headers.get('content-type'))
  75.                 if mimetype.startswith('multipart/'):
  76.                     sub_boundary = params['boundary']
  77.                     sub_parts = read_multipart(fileobj, boundary=sub_boundary)
  78.                     if boundary is not None:
  79.                         yield headers, True, sub_parts
  80.                         headers.clear()
  81.                         del buf[:]
  82.                     else:
  83.                         for part in sub_parts:
  84.                             yield part
  85.                         return
  86.  
  87.         elif line.replace(CRLF, '\n') == next_boundary:
  88.             # We've reached the start of a new part, as indicated by the
  89.             # boundary
  90.             if headers:
  91.                 if not outer:
  92.                     yield _current_part()
  93.                 else:
  94.                     outer = False
  95.                 headers.clear()
  96.                 del buf[:]
  97.             in_headers = True
  98.  
  99.         elif line.replace(CRLF, '\n') == last_boundary:
  100.             # We're done with this multipart envelope
  101.             break
  102.  
  103.         else:
  104.             buf.append(line)
  105.  
  106.     if not outer and headers:
  107.         yield _current_part()
  108.  
  109.  
  110. class MultipartWriter(object):
  111.  
  112.     def __init__(self, fileobj, headers=None, subtype='mixed', boundary=None):
  113.         self.fileobj = fileobj
  114.         if boundary is None:
  115.             boundary = self._make_boundary()
  116.         self.boundary = boundary
  117.         if headers is None:
  118.             headers = {}
  119.         headers['Content-Type'] = 'multipart/%s; boundary="%s"' % (
  120.             subtype, self.boundary
  121.         )
  122.         self._write_headers(headers)
  123.  
  124.     def open(self, headers=None, subtype='mixed', boundary=None):
  125.         self.fileobj.write('--')
  126.         self.fileobj.write(self.boundary)
  127.         self.fileobj.write(CRLF)
  128.         return MultipartWriter(self.fileobj, headers=headers, subtype=subtype,
  129.                                boundary=boundary)
  130.  
  131.     def add(self, mimetype, content, headers=None):
  132.         self.fileobj.write('--')
  133.         self.fileobj.write(self.boundary)
  134.         self.fileobj.write(CRLF)
  135.         if headers is None:
  136.             headers = {}
  137.         if isinstance(content, unicode):
  138.             ctype, params = parse_header(mimetype)
  139.             if 'charset' in params:
  140.                 content = content.encode(params['charset'])
  141.             else:
  142.                 content = content.encode('utf-8')
  143.                 mimetype = mimetype + ';charset=utf-8'
  144.         headers['Content-Type'] = mimetype
  145.         if content:
  146.             headers['Content-Length'] = str(len(content))
  147.             headers['Content-MD5'] = b64encode(md5(content).digest())
  148.         self._write_headers(headers)
  149.         if content:
  150.             # XXX: throw an exception if a boundary appears in the content??
  151.             self.fileobj.write(content)
  152.             self.fileobj.write(CRLF)
  153.  
  154.     def close(self):
  155.         self.fileobj.write('--')
  156.         self.fileobj.write(self.boundary)
  157.         self.fileobj.write('--')
  158.         self.fileobj.write(CRLF)
  159.  
  160.     def _make_boundary(self):
  161.         try:
  162.             from uuid import uuid4
  163.             return '==' + uuid4().hex + '=='
  164.         except ImportError:
  165.             from random import randrange
  166.             token = randrange(sys.maxint)
  167.             format = '%%0%dd' % len(repr(sys.maxint - 1))
  168.             return '===============' + (fmt % token) + '=='
  169.  
  170.     def _write_headers(self, headers):
  171.         if headers:
  172.             for name in sorted(headers.keys()):
  173.                 self.fileobj.write(name)
  174.                 self.fileobj.write(': ')
  175.                 self.fileobj.write(headers[name])
  176.                 self.fileobj.write(CRLF)
  177.         self.fileobj.write(CRLF)
  178.  
  179.     def __enter__(self):
  180.         return self
  181.  
  182.     def __exit__(self, exc_type, exc_val, exc_tb):
  183.         self.close()
  184.  
  185.  
  186. def write_multipart(fileobj, subtype='mixed', boundary=None):
  187.     r"""Simple streaming MIME multipart writer.
  188.  
  189.     This function returns a `MultipartWriter` object that has a few methods to
  190.     control the nested MIME parts. For example, to write a flat multipart
  191.     envelope you call the ``add(mimetype, content, [headers])`` method for
  192.     every part, and finally call the ``close()`` method.
  193.  
  194.     >>> from StringIO import StringIO
  195.  
  196.     >>> buf = StringIO()
  197.     >>> envelope = write_multipart(buf, boundary='==123456789==')
  198.     >>> envelope.add('text/plain', 'Just testing')
  199.     >>> envelope.close()
  200.     >>> print buf.getvalue().replace('\r\n', '\n')
  201.     Content-Type: multipart/mixed; boundary="==123456789=="
  202.     <BLANKLINE>
  203.     --==123456789==
  204.     Content-Length: 12
  205.     Content-MD5: nHmX4a6el41B06x2uCpglQ==
  206.     Content-Type: text/plain
  207.     <BLANKLINE>
  208.     Just testing
  209.     --==123456789==--
  210.     <BLANKLINE>
  211.  
  212.     Note that an explicit boundary is only specified for testing purposes. If
  213.     the `boundary` parameter is omitted, the multipart writer will generate a
  214.     random string for the boundary.
  215.  
  216.     To write nested structures, call the ``open([headers])`` method on the
  217.     respective envelope, and finish each envelope using the ``close()`` method:
  218.  
  219.     >>> buf = StringIO()
  220.     >>> envelope = write_multipart(buf, boundary='==123456789==')
  221.     >>> part = envelope.open(boundary='==abcdefghi==')
  222.     >>> part.add('text/plain', 'Just testing')
  223.     >>> part.close()
  224.     >>> envelope.close()
  225.     >>> print buf.getvalue().replace('\r\n', '\n') #:doctest +ELLIPSIS
  226.     Content-Type: multipart/mixed; boundary="==123456789=="
  227.     <BLANKLINE>
  228.     --==123456789==
  229.     Content-Type: multipart/mixed; boundary="==abcdefghi=="
  230.     <BLANKLINE>
  231.     --==abcdefghi==
  232.     Content-Length: 12
  233.     Content-MD5: nHmX4a6el41B06x2uCpglQ==
  234.     Content-Type: text/plain
  235.     <BLANKLINE>
  236.     Just testing
  237.     --==abcdefghi==--
  238.     --==123456789==--
  239.     <BLANKLINE>
  240.     
  241.     :param fileobj: a writable file-like object that the output should get
  242.                     written to
  243.     :param subtype: the subtype of the multipart MIME type (e.g. "mixed")
  244.     :param boundary: the boundary to use to separate the different parts
  245.     :since: 0.6
  246.     """
  247.     return MultipartWriter(fileobj, subtype=subtype, boundary=boundary)
  248.